import numpy as np
import matplotlib.pyplot as plt

from matplotlib.figure import Figure
from matplotlib.patches import ConnectionPatch
from src.common.descriptors.user_desc_lib import Software
from datetime import date
from itertools import chain

def statistics(softwares: list[Software]) -> tuple[int, dict[str, int], int, dict[str, int]]:
    not_browser = [software for software in softwares if software.platform == 'computer']
    browser = [software for software in softwares if software.platform != 'computer']
    time_not_browser = 0
    stats_not_browser = {}
    for software in not_browser:
        time_not_browser += software.time
        if software.name not in stats_not_browser:
            stats_not_browser[software.name] = software.time
        else:
            stats_not_browser[software.name] += software.time

    time_browser = 0
    stats_browser = {}
    for software in browser:
        time_browser += software.time
        try: 
            name = software.name.split('/')[2]
        except:
            name = software.name
            
        if name not in stats_browser:
            stats_browser[name] = software.time
        else:
            stats_browser[name] += software.time

    return time_browser, stats_browser, time_not_browser, stats_not_browser

def software_usage(softwares: list[Software]) -> Figure:

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
    fig.subplots_adjust(wspace=-0.5)

    # 处理数据，获取统计数据
    time_browser, stats_browser, time_not_browser, stats_not_browser = statistics(softwares)
    total_time = time_browser + time_not_browser

    # 转为比例，按使用比例降序排列
    stats_browser = [(k, v/total_time) for k, v in stats_browser.items()]
    stats_browser.sort(key=lambda x: x[1], reverse=True)
    
    stats_not_browser = [(k, v/total_time) for k, v in stats_not_browser.items()]
    stats_not_browser.sort(key=lambda x: x[1], reverse=True)

    # 软件显示使用前5名
    ratios_not_browser = [*[ele[1] for ele in stats_not_browser[0:5]]]
    # 精度导致可能减出负数
    other_not_browser = max(0, 1 - sum(ratios_not_browser) - time_browser/total_time)
    ratios_not_browser.append(other_not_browser)
    overall_ratios = [time_browser/total_time, ] + ratios_not_browser
    labels = ['浏览器', *[ele[0] for ele in stats_not_browser[0:5]], '其他']
    explode = [0.1, *[0 for _ in overall_ratios[1:]]]
    
    # 计算起始角度
    angle = -180 * overall_ratios[0]

    # 绘制饼图
    wedges, *_ = ax1.pie(overall_ratios, autopct='%1.1f%%', startangle=angle,labels=labels, explode=explode)

    # 绘制浏览器详情柱状图
    ratios_browser = [*[ele[1] for ele in stats_browser[0:9]]]
    other_browser = max(0, 1 - sum(ratios_browser))
    ratios_browser.append(other_browser)
    labels_browser = [*[ele[0] for ele in stats_browser[0:9]], '其他']
    bottom = 1
    width = .2

    # 用于计算透明度
    max_j = max(1, len(ratios_browser) - 1)
    
    for j, (height, label) in enumerate([*zip(ratios_browser, labels_browser)]):
        safe_alpha = 1 - 0.9 * (j/max_j)
        bottom -= height
        ax2.bar(0, height, width, bottom=bottom, color='C0', label=label, alpha=safe_alpha)

    # 设置右侧图表样式
    ax2.set_title('浏览器使用情况')
    ax2.legend()
    ax2.axis('off')
    ax2.set_xlim(- 2.5 * width, 2.5 * width)

    # 绘制连接线
    theta1, theta2 = max(-90,wedges[0].theta1), min(90,wedges[0].theta2)
    center, r = wedges[0].center, wedges[0].r
    bar_height = sum(ratios_browser)

    # 绘制顶部连接线
    x = r * np.cos(np.pi / 180 * theta2) + center[0]
    y = r * np.sin(np.pi / 180 * theta2) + center[1]
    con = ConnectionPatch(xyA=(-width / 2, bar_height), coordsA=ax2.transData,xyB=(x, y), coordsB=ax1.transData)
    con.set_color([0, 0, 0])
    con.set_linewidth(2)
    ax2.add_artist(con)

    # 绘制底部连接线
    x = r * np.cos(np.pi / 180 * theta1) + center[0]
    y = r * np.sin(np.pi / 180 * theta1) + center[1]
    con = ConnectionPatch(xyA=(-width / 2, 0), coordsA=ax2.transData,
                        xyB=(x, y), coordsB=ax1.transData)
    con.set_color([0, 0, 0])
    ax2.add_artist(con)
    con.set_linewidth(2)

    return fig

def software_usage_none() -> Figure:
    fig = plt.figure(figsize=(5, 2))
    ax = fig.add_subplot(111)
    ax.axis('off')
    ax.text(0.5, 0.5, '无日志记录', ha='center', va='center')
    plt.tight_layout()
    return fig

def software_usage_day(day:date, softwares: dict[date, list[Software]]) -> Figure:
    software_day = softwares.get(day)
    if software_day is None:
        return software_usage_none()
    return software_usage(software_day)

def software_usage_week(week:tuple[date, date], softwares: dict[date, list[Software]]) -> Figure:
    software_week = []
    for day, software in softwares.items():
        if week[0] <= day <= week[1]:
            software_week.extend(software)
    if not software_week:
        return software_usage_none()
    return software_usage(software_week)

def software_usage_total(softwares: dict[date, list[Software]]) -> Figure:
    return software_usage(list(chain.from_iterable(softwares.values())))